#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# cve_matcher.py

import os
import re
import json
import sqlite3
import requests
import subprocess
from packaging import version
from datetime import datetime, timedelta
import gzip
import shutil
from pathlib import Path


class CVEMatcher:
    def __init__(self, db_path="/root/web/cve.db", cache_dir=".cache"):
        """
        初始化CVE匹配器
        :param db_path: SQLite数据库路径
        :param cache_dir: 缓存目录
        """
        self.db_path = db_path
        self.cache_dir = cache_dir
        self.nvd_api_key = os.getenv("NVD_API_KEY", "")  # 从环境变量获取API密钥
        os.makedirs(self.cache_dir, exist_ok=True)
        os.makedirs(os.path.dirname(self.db_path), exist_ok=True)
        self._init_db()

        # 初始化数据库
        self._init_db()

    def _init_db(self):
        """初始化SQLite数据库"""
        with sqlite3.connect(self.db_path) as conn:
            cursor = conn.cursor()

            # 创建CVE表
            cursor.execute("""
                           CREATE TABLE IF NOT EXISTS cve
                           (
                               cve_id
                               TEXT
                               PRIMARY
                               KEY,
                               published_date
                               TEXT,
                               last_modified_date
                               TEXT,
                               description
                               TEXT,
                               severity
                               TEXT,
                               cvss_score
                               REAL,
                               affected_software
                               TEXT,
                               version_start_including
                               TEXT,
                               version_start_excluding
                               TEXT,
                               version_end_excluding
                               TEXT,
                               version_end_including
                               TEXT
                           )
                           """)

            # 创建软件包表
            cursor.execute("""
                           CREATE TABLE IF NOT EXISTS software
                           (
                               name
                               TEXT,
                               version
                               TEXT,
                               PRIMARY
                               KEY
                           (
                               name,
                               version
                           )
                               )
                           """)

            # 创建索引
            cursor.execute("""
                           CREATE INDEX IF NOT EXISTS idx_software_name ON cve(affected_software)
                           """)
            cursor.execute("""
                           CREATE INDEX IF NOT EXISTS idx_software_version ON cve(
                               affected_software,
                               version_start_including,
                               version_start_excluding,
                               version_end_excluding,
                               version_end_including
                               )
                           """)

            conn.commit()

    def _download_nvd_feeds(self, year=None):
        """
        下载NVD漏洞数据馈送
        :param year: 指定年份，None表示下载所有可用年份
        :return: 下载的文件路径列表
        """
        base_url = "https://nvd.nist.gov/feeds/json/cve/2.0/nvdcve-2.0-{}.json.gz"
        downloaded_files = []

        print(f"扫描缓存目录: {self.cache_dir}")
        json_files = []
        for f in os.listdir(self.cache_dir):
            if f.startswith("nvdcve-2.0-") and f.endswith(".json"):
                json_file = os.path.join(self.cache_dir, f)
                print(f"找到缓存文件: {json_file}")
                json_files.append(json_file)
        return json_files

        if year:
            years = [year]
        else:
            current_year = datetime.now().year
            years = range(2024, current_year + 1)

        for y in years:
            url = base_url.format(y)
            gz_file = os.path.join(self.cache_dir, f"nvdcve-2.0-{y}.json.gz")
            json_file = os.path.join(self.cache_dir, f"nvdcve-2.0-{y}.json")

            # 检查文件是否已存在且是最近30天内下载的
            if os.path.exists(json_file):
                file_time = datetime.fromtimestamp(os.path.getmtime(json_file))
                if datetime.now() - file_time < timedelta(days=30):
                    print(f"使用缓存文件: {json_file}")
                    downloaded_files.append(json_file)
                    continue

            print(f"下载: {url}")
            try:
                response = requests.get(url, stream=True, timeout=30)
                response.raise_for_status()

                with open(gz_file, 'wb') as f:
                    for chunk in response.iter_content(chunk_size=8192):
                        f.write(chunk)

                # 解压文件
                with gzip.open(gz_file, 'rb') as f_in:
                    with open(json_file, 'wb') as f_out:
                        shutil.copyfileobj(f_in, f_out)

                downloaded_files.append(json_file)
                os.remove(gz_file)
            except Exception as e:
                print(f"下载或解压 {url} 失败: {str(e)}")
                if os.path.exists(gz_file):
                    os.remove(gz_file)

        return downloaded_files

    def _parse_cve_item(self, cve_item):
        """
        解析单个CVE条目(NVD 2.0格式)
        :param cve_item: CVE JSON条目
        :return: 解析后的CVE字典
        """

        try:
            # 尝试获取CVE ID用于错误报告
            cve = cve_item.get("cve", {})
            cve_id = cve.get("id")

            if not cve_id:
                return None

            en_desc = next(
                (desc["value"] for desc in cve.get("descriptions", [])
                 if desc.get("lang") == "en"),
                ""
            )

            # 获取CVSS评分
            cvss_data = cve.get("metrics", {}).get("cvssMetricV31", [{}])[0].get("cvssData", {})

            return {
                "cve_id": cve_id,
                "published_date": cve.get("published"),
                "last_modified_date": cve.get("lastModified"),
                "description": en_desc,
                "severity": cvss_data.get("baseSeverity", "UNKNOWN"),
                "cvss_score": float(cvss_data.get("baseScore", 0.0)),
                "affected_software": "",
                "version_ranges": self._parse_version_ranges(cve)
            }

        except Exception as e:
            print(f"解析 {cve.get('id', 'UNKNOWN')} 失败: {str(e)}")
            return None

    def _parse_version_ranges(self, cve_data):
        """解析受影响软件的版本范围"""
        version_ranges = []
        for config in cve_data.get("configurations", []):
            for node in config.get("nodes", []):
                for match in node.get("cpeMatch", []):
                    if self._is_web_httpd(match.get("criteria", "")):
                        version_ranges.append({
                            "software": "httpd",
                            "version_start_including": match.get("versionStartIncluding"),
                            "version_start_excluding": match.get("versionStartExcluding"),
                            "version_end_excluding": match.get("versionEndExcluding"),
                            "version_end_including": match.get("versionEndIncluding")
                        })
        return version_ranges

    def _is_web_httpd(self, cpe_uri):
        """判断是否为Web HTTP Server的CPE"""
        cpe = cpe_uri.lower()
        return (
                ("apache" in cpe and ("httpd" in cpe or "http_server" in cpe)) and
                ":a:" in cpe  # 仅匹配软件
        )

    def check_database_connection(self):
        """检查CVE数据库连接状态"""
        try:
            with sqlite3.connect(self.db_path) as conn:
                cursor = conn.cursor()

                # 检查表是否存在
                cursor.execute("SELECT name FROM sqlite_master WHERE type='table' AND name='cve';")
                table_exists = cursor.fetchone()

                # 检查记录数量
                cursor.execute("SELECT COUNT(*) FROM cve;")
                record_count = cursor.fetchone()[0]

                return {
                    "status": "connected",
                    "database_file": self.db_path,
                    "tables_exist": bool(table_exists),
                    "record_count": record_count,
                    "file_size": f"{os.path.getsize(self.db_path) / 1024 / 1024:.2f} MB"
                }
        except Exception as e:
            return {
                "status": "disconnected",
                "error": str(e)
            }

    def _import_cve_data(self, json_files):
        """
        将CVE数据导入数据库(2.0版)
        :param json_files: JSON文件路径列表
        """
        with sqlite3.connect(self.db_path) as conn:
            cursor = conn.cursor()

            print(f"准备导入 {len(json_files)} 个文件")
            for json_file in json_files:
                vulns = []

                try:
                    with open(json_file, 'r', encoding='utf-8') as f:
                        data = json.load(f)
                    vulns = data.get("vulnerabilities", []) or data.get("CVE_Items", [])
                    print(f"文件 {json_file} 包含 {len(vulns)} 条漏洞记录")

                    success_count = 0
                    for vuln in vulns:
                        cve_data = self._parse_cve_item(vuln)
                        if not cve_data or cve_data.get("cve_id") == "UNKNOWN_ID":
                            continue

                        # 插入CVE基本信息
                        cursor.execute("""
                            INSERT OR REPLACE INTO cve (
                                cve_id, published_date, last_modified_date, 
                                description, severity, cvss_score, affected_software,
                                version_start_including, version_start_excluding, 
                                version_end_excluding, version_end_including
                            ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
                        """, (
                            cve_data["cve_id"],
                            cve_data["published_date"],
                            cve_data["last_modified_date"],
                            cve_data["description"],
                            cve_data["severity"],
                            cve_data["cvss_score"],
                            cve_data["affected_software"],
                            None, None, None, None  # 版本范围在单独的表中处理
                        ))
                        conn.commit()

                        # 插入版本范围信息
                        for vr in cve_data.get("version_ranges", []):
                            cursor.execute("""
                                INSERT OR REPLACE INTO cve (
                                    cve_id, published_date, last_modified_date, 
                                    description, severity, cvss_score, affected_software,
                                    version_start_including, version_start_excluding, 
                                    version_end_excluding, version_end_including
                                ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
                            """, (
                                cve_data["cve_id"],
                                cve_data["published_date"],
                                cve_data["last_modified_date"],
                                cve_data["description"],
                                cve_data["severity"],
                                cve_data["cvss_score"],
                                vr["software"],
                                vr["version_start_including"],
                                vr["version_start_excluding"],
                                vr["version_end_excluding"],
                                vr["version_end_including"]
                            ))
                            success_count += 1

                    conn.commit()
                    print(f"成功导入 {success_count}/{len(vulns)} 条记录")

                except json.JSONDecodeError as e:
                    conn.rollback()
                    print(f"JSON解析失败 {json_file}: {str(e)}")
                except Exception as e:
                    conn.rollback()
                    print(f"导入失败 {json_file}: {str(e)}")

                    if vulns:
                        for i, sample in enumerate(vulns[:3], 1):
                            cve_id = sample.get("cve", {}).get("CVE_data_meta", {}).get("ID", "UNKNOWN_ID")
                            print(f"样本{i} CVE-ID: {cve_id}")

    def update_cve_database(self, full_update=False):
        """
        更新CVE数据库
        :param full_update: 是否执行完整更新(False则只更新最近修改的CVE)
        """
        if full_update:
            print("执行完整CVE数据库更新...")
            json_files = self._download_nvd_feeds()
            self._import_cve_data(json_files)
        else:
            print("执行增量CVE数据库更新...")
            # 这里可以添加增量更新逻辑，限于篇幅简化实现
            # 实际应用中可以通过NVD的API获取最近修改的CVE
            json_files = self._download_nvd_feeds(datetime.now().year)
            self._import_cve_data(json_files)
        print("CVE数据库更新完成")

    def query_cve_by_software(self, software_name, software_version):
        """
        查询指定软件版本的CVE漏洞
        :param software_name: 软件名称(如"httpd")
        :param software_version: 软件版本(如"2.4.6")
        :return: 匹配的CVE列表
        """
        try:
            ver = version.parse(software_version)
        except version.InvalidVersion:
            print(f"警告: 无效的版本号格式 '{software_version}'")
            return []

        with sqlite3.connect(self.db_path) as conn:
            conn.row_factory = sqlite3.Row
            cursor = conn.cursor()

            # 查询匹配的CVE
            cursor.execute("""
                           SELECT DISTINCT cve_id,
                                           description,
                                           severity,
                                           cvss_score,
                                           version_start_including,
                                           version_start_excluding,
                                           version_end_excluding,
                                           version_end_including
                           FROM cve
                           WHERE affected_software LIKE '%' || ? || '%'
                           ORDER BY cvss_score DESC
                           """, (software_name,))

            vulnerabilities = []
            for row in cursor.fetchall():
                matched = True
                try:
                    if row["version_start_including"]:
                        if ver < version.parse(row["version_start_including"]):
                            matched = False
                    if row["version_start_excluding"]:
                        if ver <= version.parse(row["version_start_excluding"]):
                            matched = False
                    if row["version_end_excluding"]:
                        if ver >= version.parse(row["version_end_excluding"]):
                            matched = False
                    if row["version_end_including"]:
                        if ver > version.parse(row["version_end_including"]):
                            matched = False
                except version.InvalidVersion:
                    print(f"警告: 无效的版本范围格式 - CVE: {row['cve_id']}")
                    continue

                if matched:
                    vulnerabilities.append(dict(row))

            return vulnerabilities

    def get_web_vulnerabilities(self, web_version):
        """
        获取Web HTTP Server的CVE漏洞
        :param web_version: Web版本号(如"2.4.6")
        :return: 漏洞列表及修复建议
        """
        # 查询httpd相关的CVE
        cves = self.query_cve_by_software("httpd", web_version)

        if not cves:
            return {
                "status": "secure",
                "message": "未发现已知漏洞",
                "web_version": web_version,
                "scan_time": datetime.now().isoformat()
            }

        # 按严重程度分类
        critical = [cve for cve in cves if cve["severity"] == "CRITICAL"]
        high = [cve for cve in cves if cve["severity"] == "HIGH"]
        medium = [cve for cve in cves if cve["severity"] == "MEDIUM"]
        low = [cve for cve in cves if cve["severity"] == "LOW"]

        # 生成修复建议
        recommendations = []
        if critical or high:
            recommendations.append({
                "priority": "immediate",
                "action": "upgrade",
                "description": f"发现 {len(critical)} 个严重漏洞和 {len(high)} 个高危漏洞",
                "steps": [
                    "备份当前配置: cp -r /etc/httpd /etc/httpd_backup",
                    "停止Web服务: systemctl stop httpd",
                    "升级软件包: yum update httpd",
                    "验证配置: webctl configtest",
                    "重启服务: systemctl start httpd"
                ]
            })

        # 检查是否有可通过配置缓解的漏洞
        config_fixes = []
        for cve in critical + high + medium:
            if "mod_ssl" in cve["description"]:
                config_fixes.append({
                    "cve": cve["cve_id"],
                    "description": cve["description"],
                    "config_change": "SSLProtocol TLSv1.2 TLSv1.3",
                    "file": "/etc/httpd/conf.d/ssl.conf"
                })

        if config_fixes:
            recommendations.append({
                "priority": "high",
                "action": "config_update",
                "description": "配置更改可缓解部分漏洞",
                "changes": config_fixes
            })

        return {
            "status": "vulnerable",
            "web_version": web_version,
            "scan_time": datetime.now().isoformat(),
            "critical_count": len(critical),
            "high_count": len(high),
            "medium_count": len(medium),
            "low_count": len(low),
            "vulnerabilities": cves,
            "recommendations": recommendations
        }

    def get_software_version(self, software_name):
        """
        获取已安装软件的版本号
        :param software_name: 软件名称(如"httpd")
        :return: 版本号字符串
        """
        try:
            if software_name == "httpd":
                # 获取Web版本
                try:
                    output = subprocess.check_output(["httpd", "-v"], stderr=subprocess.STDOUT).decode()
                    # 尝试多种匹配模式
                    patterns = [
                        r"Apache/([\d.]+)",
                        r"Server version: Apache/([\d.]+)",
                        r"Apache HTTP Server/([\d.]+)"
                    ]
                    for pattern in patterns:
                        match = re.search(pattern, output)
                        if match:
                            return match.group(1)
                except Exception:
                    pass

                # 尝试通过rpm查询
                try:
                    output = subprocess.check_output(["rpm", "-q", "httpd", "--queryformat", "%{VERSION}"])
                    return output.decode().strip()
                except Exception:
                    pass

                # 尝试通过配置文件查找
                try:
                    with open("/etc/httpd/conf/httpd.conf", "r") as f:
                        for line in f:
                            if "ServerVersion" in line:
                                return line.split()[-1]
                except Exception:
                    pass

                return "unknown"
        except Exception as e:
            print(f"获取软件版本时出错: {str(e)}")
            return "unknown"

    # 添加方法来检查数据库中的CVE记录
    def check_cve_records(self, cve_id):
        with sqlite3.connect(self.db_path) as conn:
            conn.row_factory = sqlite3.Row
            cursor = conn.cursor()
            cursor.execute("SELECT * FROM cve WHERE cve_id = ?", (cve_id,))
            return cursor.fetchall()


if __name__ == "__main__":
    # 使用示例
    matcher = CVEMatcher(cache_dir="/root/web/.cache")

    # 检查数据库连接状态
    db_status = matcher.check_database_connection()
    print("\nCVE数据库连接状态:")
    print(json.dumps(db_status, indent=2))

    if db_status["status"] != "connected" or db_status["record_count"] == 0:
        print("\n数据库未初始化或为空，正在更新CVE数据库...")
        matcher.update_cve_database(full_update=True)

    # 更新CVE数据库(首次使用需要完整更新)
    if not os.path.exists(matcher.db_path):
        print("首次运行，正在下载完整的CVE数据库...")
        matcher.update_cve_database(full_update=True)
    else:
        # 每周执行一次完整更新
        db_age = datetime.now() - datetime.fromtimestamp(os.path.getmtime(matcher.db_path))
        if db_age > timedelta(days=7):
            print("CVE数据库超过7天未更新，正在更新...")
            matcher.update_cve_database(full_update=True)

    # 获取Web版本并检查漏洞
    web_version = matcher.get_software_version("httpd")
    print(f"检测到Web版本: {web_version}")

    if web_version == "unknown":
        print("无法自动确定Web版本，请尝试以下方法手动检查:")
        print("1. 运行命令: httpd -v 或 web2 -v")
        print("2. 运行命令: rpm -q httpd --queryformat '%{VERSION}'")
        print("3. 检查 /etc/httpd/conf/httpd.conf 文件中的版本信息")

        # 允许手动输入版本号
        manual_version = input("请输入Web版本号(如2.4.6，直接回车跳过): ").strip()
        if manual_version:
            cve_result = matcher.get_web_vulnerabilities(manual_version)
            print(json.dumps(cve_result, indent=2, ensure_ascii=False))

    else:
        cve_result = matcher.get_web_vulnerabilities(web_version)
        print(json.dumps(cve_result, indent=2, ensure_ascii=False))

    # 保存报告到文件
    # 确保web_version目录存在
    web_version_dir = os.path.join(os.path.dirname(__file__), "web_version")
    os.makedirs(web_version_dir, exist_ok=True)
    
    # 构建报告文件路径
    report_file = f"web_cve_report_{datetime.now().strftime('%Y%m%d_%H%M%S')}.json"
    report_path = os.path.join(web_version_dir, report_file)
    
    # 保存报告
    with open(report_path, 'w') as f:
        json.dump(cve_result, f, indent=2)
    print(f"报告已保存到 {report_path}")